"""
向量化服务
"""
import numpy as np
from typing import List, Dict, Any, Optional, Union
from langchain.embeddings.base import Embeddings
from langchain_community.embeddings import (
    HuggingFaceEmbeddings,
    OpenAIEmbeddings,
    SentenceTransformerEmbeddings,
)

from ..config.settings import settings


def normalize(embeddings: List[List[float]]) -> List[List[float]]:
    """L2范数归一化"""
    embeddings = np.array(embeddings)
    norms = np.linalg.norm(embeddings, axis=1, keepdims=True)
    norms = np.where(norms == 0, 1, norms)  # 避免除零
    normalized_embeddings = embeddings / norms
    return normalized_embeddings.tolist()


class EmbeddingService:
    """向量化服务"""
    
    def __init__(self, model_name: str = None, **kwargs):
        self.model_name = model_name or settings.embedding.model_name
        self.device = kwargs.get("device", settings.embedding.device)
        self.normalize = kwargs.get("normalize", settings.embedding.normalize)
        self.batch_size = kwargs.get("batch_size", settings.embedding.batch_size)
        
        self.embedding_model = self._create_embedding_model()
    
    def _create_embedding_model(self) -> Embeddings:
        """创建嵌入模型"""
        try:
            if self.model_name == "text-embedding-ada-002":
                # OpenAI嵌入模型
                return OpenAIEmbeddings(
                    model=self.model_name,
                    openai_api_key=settings.get("openai_api_key", ""),
                )
            elif "bge-" in self.model_name:
                # BGE嵌入模型
                query_instruction = ""
                if "zh" in self.model_name:
                    query_instruction = "为这个句子生成表示以用于检索相关文章："
                elif "en" in self.model_name:
                    query_instruction = "Represent this sentence for searching relevant passages:"
                
                return HuggingFaceEmbeddings(
                    model_name=self.model_name,
                    model_kwargs={"device": self.device},
                    encode_kwargs={"normalize_embeddings": self.normalize},
                    query_instruction=query_instruction,
                )
            else:
                # 默认使用SentenceTransformer
                return SentenceTransformerEmbeddings(
                    model_name=self.model_name,
                    model_kwargs={"device": self.device},
                    encode_kwargs={"normalize_embeddings": self.normalize},
                )
        except Exception as e:
            print(f"创建嵌入模型失败: {e}")
            # 回退到默认模型
            return SentenceTransformerEmbeddings(
                model_name="all-MiniLM-L6-v2",
                model_kwargs={"device": self.device},
                encode_kwargs={"normalize_embeddings": self.normalize},
            )
    
    def embed_documents(self, texts: List[str]) -> List[List[float]]:
        """批量向量化文档"""
        if not texts:
            return []
        
        try:
            embeddings = self.embedding_model.embed_documents(texts)
            
            if self.normalize:
                embeddings = normalize(embeddings)
            
            return embeddings
        except Exception as e:
            print(f"向量化文档失败: {e}")
            return []
    
    def embed_query(self, text: str) -> List[float]:
        """向量化查询文本"""
        if not text:
            return []
        
        try:
            embedding = self.embedding_model.embed_query(text)
            
            if self.normalize:
                # 转换为2D数组进行归一化
                embedding_2d = np.array([embedding])
                normalized = normalize(embedding_2d)
                return normalized[0]
            
            return embedding
        except Exception as e:
            print(f"向量化查询失败: {e}")
            return []
    
    async def aembed_documents(self, texts: List[str]) -> List[List[float]]:
        """异步批量向量化文档"""
        if not texts:
            return []
        
        try:
            embeddings = await self.embedding_model.aembed_documents(texts)
            
            if self.normalize:
                embeddings = normalize(embeddings)
            
            return embeddings
        except Exception as e:
            print(f"异步向量化文档失败: {e}")
            return []
    
    async def aembed_query(self, text: str) -> List[float]:
        """异步向量化查询文本"""
        if not text:
            return []
        
        try:
            embedding = await self.embedding_model.aembed_query(text)
            
            if self.normalize:
                # 转换为2D数组进行归一化
                embedding_2d = np.array([embedding])
                normalized = normalize(embedding_2d)
                return normalized[0]
            
            return embedding
        except Exception as e:
            print(f"异步向量化查询失败: {e}")
            return []
    
    def get_embedding_dimension(self) -> int:
        """获取向量维度"""
        try:
            # 使用一个简单的测试文本获取维度
            test_embedding = self.embed_query("test")
            return len(test_embedding)
        except:
            return 768  # 默认维度


class EmbeddingAdapter:
    """嵌入适配器，兼容不同接口"""
    
    def __init__(self, model_name: str = None, **kwargs):
        self.embedding_service = EmbeddingService(model_name, **kwargs)
    
    def embed_documents(self, texts: List[str]) -> List[List[float]]:
        """批量向量化文档"""
        return self.embedding_service.embed_documents(texts)
    
    def embed_query(self, text: str) -> List[float]:
        """向量化查询文本"""
        return self.embedding_service.embed_query(text)
    
    async def aembed_documents(self, texts: List[str]) -> List[List[float]]:
        """异步批量向量化文档"""
        return await self.embedding_service.aembed_documents(texts)
    
    async def aembed_query(self, text: str) -> List[float]:
        """异步向量化查询文本"""
        return await self.embedding_service.aembed_query(text) 